import pickle
import os
import numpy as np
from python_ai.common.xcommon import *
from python_ai.category.NumpyNet.utils import *
import sys
import cv2 as cv

sep('Load data')
BASE_DIR, FILE_NAME = os.path.split(__file__)
save_path = os.path.join(BASE_DIR, '_data', 'mnist_dataset_gen.py.pkl')
if not os.path.exists(save_path):
    raise Exception(f'{save_path} does not exist. Please run mnist_dataset_gen.py to generate it.')
print(f'Loading data from {save_path}')
with open(save_path, 'rb') as f:
    data_dict = pickle.load(f)

sep('Write to files')
OVERWRITE = 0
BASE_DIR, FILE_NAME = os.path.split(__file__)
xdir = os.path.join(BASE_DIR, '_data', FILE_NAME)
os.makedirs(xdir, exist_ok=True)

for xtype in ('train', 'val', 'test'):
    print(f'Started processing type {xtype}')
    xpath = os.path.join(xdir, xtype)
    os.makedirs(xpath, exist_ok=True)
    x = data_dict[f'x_{xtype}']
    y = data_dict[f'y_{xtype}']
    for i, img in enumerate(x):
        # if i == 3:
        #     break
        # print(img.shape)
        no = i + 1
        if no % 2000 == 0:
            print(f'{no}')
        gt = y[i]
        filename = os.path.join(xpath, f'{no}_{gt}.png')
        if not OVERWRITE and os.path.exists(filename):
            continue
        # print(img.shape, img.min(), img.max(), img.dtype)
        cv.imwrite(filename, img)
    print(f'Finished processing type {xtype}: {no}')
